import os
from xml.dom.minidom import Document
from xml.dom.minidom import parse
import xml.dom.minidom
import numpy as np
import csv
import cv2
from tqdm import tqdm
import json


def WriterXMLFiles(filename, path, box_list, label_list, w, h, d):

    # dict_box[filename]=json_dict[filename]
    doc = xml.dom.minidom.Document()
    root = doc.createElement('annotation')
    doc.appendChild(root)

    foldername = doc.createElement("folder")
    foldername.appendChild(doc.createTextNode("JPEGImages"))
    root.appendChild(foldername)

    nodeFilename = doc.createElement('filename')
    nodeFilename.appendChild(doc.createTextNode(filename))
    root.appendChild(nodeFilename)

    pathname = doc.createElement("path")
    pathname.appendChild(doc.createTextNode("xxxx"))
    root.appendChild(pathname)

    sourcename=doc.createElement("source")

    databasename = doc.createElement("database")
    databasename.appendChild(doc.createTextNode("Unknown"))
    sourcename.appendChild(databasename)

    annotationname = doc.createElement("annotation")
    annotationname.appendChild(doc.createTextNode("xxx"))
    sourcename.appendChild(annotationname)

    imagename = doc.createElement("image")
    imagename.appendChild(doc.createTextNode("xxx"))
    sourcename.appendChild(imagename)

    flickridname = doc.createElement("flickrid")
    flickridname.appendChild(doc.createTextNode("0"))
    sourcename.appendChild(flickridname)

    root.appendChild(sourcename)

    nodesize = doc.createElement('size')
    nodewidth = doc.createElement('width')
    nodewidth.appendChild(doc.createTextNode(str(w)))
    nodesize.appendChild(nodewidth)
    nodeheight = doc.createElement('height')
    nodeheight.appendChild(doc.createTextNode(str(h)))
    nodesize.appendChild(nodeheight)
    nodedepth = doc.createElement('depth')
    nodedepth.appendChild(doc.createTextNode(str(d)))
    nodesize.appendChild(nodedepth)
    root.appendChild(nodesize)

    segname = doc.createElement("segmented")
    segname.appendChild(doc.createTextNode("0"))
    root.appendChild(segname)

    for (box, label) in zip(box_list, label_list):

        nodeobject = doc.createElement('object')
        nodename = doc.createElement('name')
        nodename.appendChild(doc.createTextNode(str(label)))
        nodeobject.appendChild(nodename)
        nodedifficult = doc.createElement('difficult')
        nodedifficult.appendChild(doc.createTextNode("0"))
        nodeobject.appendChild(nodedifficult)
        nodebndbox = doc.createElement('bndbox')
        nodex1 = doc.createElement('xmin')
        nodex1.appendChild(doc.createTextNode(str(box[0])))
        nodebndbox.appendChild(nodex1)
        nodey1 = doc.createElement('ymin')
        nodey1.appendChild(doc.createTextNode(str(box[1])))
        nodebndbox.appendChild(nodey1)
        nodex2 = doc.createElement('xmax')
        nodex2.appendChild(doc.createTextNode(str(box[0] + box[2])))
        nodebndbox.appendChild(nodex2)
        nodey2 = doc.createElement('ymax')
        nodey2.appendChild(doc.createTextNode(str(box[1] + box[3])))
        nodebndbox.appendChild(nodey2)

        # ang = doc.createElement('angle')
        # ang.appendChild(doc.createTextNode(str(angle)))
        # nodebndbox.appendChild(ang)
        nodeobject.appendChild(nodebndbox)
        root.appendChild(nodeobject)
    fp = open(path + filename, 'w')
    doc.writexml(fp, indent='\n')
    fp.close()


def load_annoataion(p):
    '''
    load annotation from the text file
    :param p:
    :return:
    '''
    text_polys = []
    text_tags = []
    if not os.path.exists(p):
        return np.array(text_polys, dtype=np.float32)
    with open(p, 'r') as f:
        reader = csv.reader(f)
        for line in reader:
            label = 'text'
            # strip BOM. \ufeff for python3,  \xef\xbb\bf for python2
            line = [i.strip('\ufeff').strip('\xef\xbb\xbf') for i in line]

            x1, y1, x2, y2, x3, y3, x4, y4 = list(map(float, line[:8]))
            text_polys.append([x1, y1, x2, y2, x3, y3, x4, y4])
            text_tags.append(label)

        return np.array(text_polys, dtype=np.int32), np.array(text_tags, dtype=np.str)

if __name__ == "__main__":
    data_root = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev'
    # txt_path = '/media/hp208/4t/zhaoxingjie/RS/data/rssrai2019_object_detection/val/labeltxt/'
    xml_path = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev-voc-format/VOC2007/Annotations/'
    img_path = '/media/hp208/4t/zhaoxingjie/project_graduation/data/test-dev-voc-format/VOC2007/JPEGImages/'
    # print(os.path.exists(data_root))
    vid_names = os.listdir(data_root)
    vid_names.sort()
    for count, vid_name in enumerate(tqdm(vid_names)):
        img_idx = 1             # img_name
        json_path = os.path.join(data_root, vid_name, 'RGB_label.json')
        with open(json_path, 'r') as f:
            gts = json.load(f)  # IR ---> RGB
            exist = gts['exist']
            gt_rect = gts['gt_rect']
        video_cap = cv2.VideoCapture(os.path.join(data_root, vid_name, 'RGB.mp4'))
        for rect in gt_rect:
            img_name = os.path.join(img_path, vid_name + '_{:05d}.jpg'.format(img_idx))
            xml_name = os.path.join(xml_path, vid_name + '_{:05d}.xml'.format(img_idx))
            # print(img_name)
            # print(xml_name)
            ref, img = video_cap.read()
            if not ref:
                break
            if len(rect)==0:
                img_idx += 1
                continue

            h, w, d = img.shape
            WriterXMLFiles(xml_name.split('/')[-1], xml_path, [rect], ['uav'], w, h, d)
            cv2.imwrite(img_name, img)
            img_idx += 1
        video_cap.release()